import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt
import numpy as np


def getTrainSetAndTestSet(DataPath):
    data = pd.read_csv(DataPath)  # 读取CSV文件
    X = data[['AT', 'V', 'AP', 'RH']]  # 选取AT, V, AP, RH作为特征
    y = data[['PE']]  # 选取PE作为标签
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1)  # 随机划分训练集和测试集，默认25%作为测试集
    print("训练集特征维度:", X_train.shape)
    print("训练集标签维度:", y_train.shape)
    print("测试集特征维度:", X_test.shape)
    print("测试集标签维度:", y_test.shape)
    return X_train, X_test, y_train, y_test


def TrainLinearRegression(X_train, y_train):
    linreg = LinearRegression()  # 创建线性回归模型
    linreg.fit(X_train, y_train)  # 使用训练集数据训练模型
    print("线性回归截距:", linreg.intercept_)
    print("线性回归系数:", linreg.coef_)
    return linreg


def EvaluationModel(linreg, X_test, y_test):
    y_pred = linreg.predict(X_test)  # 使用模型预测测试集
    # 修正：将DataFrame转换为numpy数组进行计算
    mse = np.mean((y_pred - y_test.values) ** 2)  # 计算均方误差
    print("均方误差MSE:", mse)
    rmse = np.sqrt(mse)  # 计算均方根误差
    print("均方根误差RMSE:", rmse)
    return y_pred


def Visualization(y_test, y_pred):
    # 修正：将DataFrame转换为numpy数组进行绘图
    y_test_array = y_test.values.flatten()
    y_pred_array = y_pred.flatten()

    fig, ax = plt.subplots(figsize=(10, 6))  # 创建图形和轴
    ax.scatter(y_test_array, y_pred_array, alpha=0.5)  # 绘制散点图，添加透明度
    ax.plot([y_test_array.min(), y_test_array.max()],
            [y_test_array.min(), y_test_array.max()],
            'k--', lw=2)  # 绘制对角线
    ax.set_xlabel("Measured", fontsize=12)  # 设置x轴标签
    ax.set_ylabel("Predicted", fontsize=12)  # 设置y轴标签
    ax.set_title("实际值 vs 预测值", fontsize=14)  # 添加标题
    ax.grid(True, linestyle='--', alpha=0.7)  # 添加网格线

    # 修正：根据运行环境选择合适的显示方式
    if '__file__' not in globals():  # 如果在Jupyter环境中
        plt.show()
    else:  # 如果在脚本环境中
        plt.savefig('prediction_vs_actual.png', dpi=300, bbox_inches='tight')
        print("图形已保存为 prediction_vs_actual.png")

    plt.close(fig)  # 关闭图形


if __name__ == "__main__":
    data_path = "Folds5x2_pp.csv"  # 数据集路径
    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(data_path)  # 获取训练集和测试集
    linreg_model = TrainLinearRegression(X_train, y_train)  # 训练线性回归模型
    y_pred = EvaluationModel(linreg_model, X_test, y_test)  # 评估模型性能
    Visualization(y_test, y_pred)  # 可视化结果